#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
攻击链数据接收服务器
监听 http://192.168.200.53:3620/api/attack/chain 端点
接收Vue组件发送的CSV格式攻击链数据
"""

import json
import csv
import io
import os
import datetime
from flask import Flask, request, jsonify
from flask_cors import CORS

app = Flask(__name__)
CORS(app)  # 允许跨域请求

# 数据存储目录
DATA_DIR = "received_data"
if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)

def save_received_data(data_type, content, headers=None):
    """保存接收到的数据到文件"""
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # 保存原始数据
    filename = f"{data_type}_{timestamp}.txt"
    filepath = os.path.join(DATA_DIR, filename)
    
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(f"接收时间: {datetime.datetime.now().isoformat()}\n")
        f.write(f"数据类型: {data_type}\n")
        f.write(f"Content-Type: {headers.get('Content-Type', 'unknown') if headers else 'unknown'}\n")
        f.write("-" * 50 + "\n")
        f.write(content)
    
    print(f"✅ 数据已保存到: {filepath}")
    return filepath

def parse_csv_data(csv_content):
    """解析CSV数据并返回结构化信息"""
    try:
        csv_reader = csv.DictReader(io.StringIO(csv_content))
        rows = list(csv_reader)
        
        if not rows:
            return {"valid": False, "error": "CSV数据为空"}
        
        # 验证CSV格式
        required_columns = ['parent_drs_in_ip', 'parent_drs_backend_ip', 'child_drs_in_ips']
        if not all(col in rows[0] for col in required_columns):
            return {"valid": False, "error": f"缺少必需的列: {required_columns}"}
        
        # 分析数据
        analysis = {
            "valid": True,
            "row_count": len(rows),
            "columns": list(rows[0].keys()),
            "sample_data": rows[:3] if len(rows) > 3 else rows,
            "attack_chains": []
        }
        
        # 提取攻击链信息
        for i, row in enumerate(rows):
            chain_info = {
                "chain_id": i + 1,
                "parent_frontend": row['parent_drs_in_ip'],
                "parent_backend": row['parent_drs_backend_ip'],
                "child_frontend": row['child_drs_in_ips'].strip('"').split(',') if row['child_drs_in_ips'] else []
            }
            analysis["attack_chains"].append(chain_info)
        
        return analysis
        
    except Exception as e:
        return {"valid": False, "error": f"CSV解析失败: {str(e)}"}

@app.route('/api/attack/chain', methods=['POST'])
def receive_attack_chain():
    """接收攻击链数据的主端点"""
    try:
        # 获取请求信息
        content_type = request.headers.get('Content-Type', '')
        content_length = request.headers.get('Content-Length', '0')
        
        print(f"\n{'='*60}")
        print(f"🔄 接收到新的攻击链数据")
        print(f"时间: {datetime.datetime.now().isoformat()}")
        print(f"Content-Type: {content_type}")
        print(f"Content-Length: {content_length}")
        print(f"客户端IP: {request.remote_addr}")
        print(f"{'='*60}")
        
        # 获取请求体
        if request.is_json:
            # 如果是JSON格式
            data = request.get_json()
            content = json.dumps(data, ensure_ascii=False, indent=2)
            data_type = "json"
        else:
            # 如果是纯文本格式（CSV）
            content = request.get_data(as_text=True)
            data_type = "csv"
        
        print(f"📄 接收到的数据内容:")
        print(content)
        print(f"\n📊 数据长度: {len(content)} 字符")
        
        # 保存原始数据
        save_received_data(data_type, content, request.headers)
        
        # 如果是CSV数据，进行解析
        if data_type == "csv":
            csv_analysis = parse_csv_data(content)
            if csv_analysis["valid"]:
                print(f"\n✅ CSV数据解析成功:")
                print(f"   - 行数: {csv_analysis['row_count']}")
                print(f"   - 列数: {len(csv_analysis['columns'])}")
                print(f"   - 攻击链数量: {len(csv_analysis['attack_chains'])}")
                
                # 保存解析后的JSON数据
                json_filename = f"parsed_attack_chain_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
                json_filepath = os.path.join(DATA_DIR, json_filename)
                with open(json_filepath, 'w', encoding='utf-8') as f:
                    json.dump(csv_analysis, f, ensure_ascii=False, indent=2)
                print(f"✅ 解析数据已保存到: {json_filepath}")
                
                # 显示攻击链详情
                print(f"\n🔗 攻击链详情:")
                for chain in csv_analysis["attack_chains"]:
                    print(f"   链 {chain['chain_id']}: {chain['parent_frontend']} -> {chain['parent_backend']} -> {chain['child_frontend']}")
            else:
                print(f"❌ CSV数据解析失败: {csv_analysis['error']}")
        
        # 返回成功响应
        response_data = {
            "success": True,
            "message": "请求已接收，攻击链正在后台部署。请在服务器控制台查看详细日志。",
            "timestamp": datetime.datetime.now().isoformat(),
            "data_type": data_type,
            "content_length": len(content)
        }
        
        if data_type == "csv":
            response_data["csv_analysis"] = csv_analysis
        
        print(f"\n✅ 响应发送成功")
        return jsonify(response_data), 202
        
    except Exception as e:
        error_msg = f"处理请求时发生错误: {str(e)}"
        print(f"\n❌ {error_msg}")
        return jsonify({
            "success": False,
            "error": error_msg,
            "timestamp": datetime.datetime.now().isoformat()
        }), 500

@app.route('/api/attack/status', methods=['GET'])
def get_attack_status():
    """获取攻击状态"""
    try:
        # 统计接收到的数据文件
        files = [f for f in os.listdir(DATA_DIR) if f.endswith('.txt')]
        
        status = {
            "server_status": "running",
            "total_received": len(files),
            "last_update": datetime.datetime.now().isoformat(),
            "data_directory": DATA_DIR
        }
        
        return jsonify(status), 200
        
    except Exception as e:
        return jsonify({
            "success": False,
            "error": f"获取状态失败: {str(e)}"
        }), 500

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    return jsonify({
        "status": "healthy",
        "timestamp": datetime.datetime.now().isoformat(),
        "message": "攻击链数据接收服务器运行正常"
    }), 200

@app.route('/', methods=['GET'])
def index():
    """根路径信息"""
    return jsonify({
        "service": "攻击链数据接收服务器",
        "version": "1.0.0",
        "endpoints": {
            "POST /api/attack/chain": "接收攻击链数据",
            "GET /api/attack/status": "获取攻击状态",
            "GET /health": "健康检查"
        },
        "data_directory": DATA_DIR,
        "timestamp": datetime.datetime.now().isoformat()
    }), 200

if __name__ == '__main__':
    print(f"""
{'='*60}
🚀 攻击链数据接收服务器启动中...
{'='*60}
📍 服务地址: http://192.168.200.53:3620
📁 数据存储目录: {os.path.abspath(DATA_DIR)}
🔄 支持的端点:
   - POST /api/attack/chain (接收攻击链数据)
   - GET /api/attack/status (获取状态)
   - GET /health (健康检查)
   - GET / (服务信息)

📊 监听中，等待Vue组件发送数据...
{'='*60}
    """)
    
    # 启动服务器
    app.run(
        host='0.0.0.0',
        port=3620,
        debug=True,
        threaded=True
    )




